###
###
###

class Trainer:
    def __init__(self, model, optimizer, criterion, device):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.device = device

    def train_step(self, inputs, labels):
        inputs = inputs.to(self.device)
        labels = labels.to(self.device)

        self.optimizer.zero_grad()

        outputs = self.model(inputs)
        loss = self.criterion(outputs, labels)

        loss.backward()
        self.optimizer.step()

        return loss.item()

    # def eval_step(self, inputs, labels):
    #     inputs = inputs.to(self.device)
    #     labels = labels.to(self.device)

    #     with torch.no_grad():
    #         outputs = self.model(inputs)
    #         loss = self.criterion(outputs, labels)

    #     return loss.item() 
    